import os
import pandas as pd
import numpy as np
import random
import openpyxl
from openpyxl.styles import Alignment, Font, PatternFill, Border, Side
from openpyxl.utils import get_column_letter


def sample_meter_list(
    input_dir="input",
    output_dir="output",
    file_name="用户清单.xlsx",
    header_name=None,
    filter_values="",
    total_count=400,
    output_file="抽查结果.xlsx",
    columns_to_keep=None,
):
    """
    从用户清单中随机抽查数据

    参数:
    - input_dir: 输入目录
    - output_dir: 输出目录
    - file_name: 输入文件名
    - header_name: 筛选表头名称（可选）
    - filter_values: 筛选值（逗号分隔，可选）
    - total_count: 需要筛选的总条数
    - output_file: 输出文件名
    - columns_to_keep: 需要保留的列（列表，可选）

    返回:
    - 输出文件路径
    """
    # 文件路径
    input_path = os.path.join(input_dir, file_name)

    # 检查文件是否存在
    if not os.path.exists(input_path):
        raise FileNotFoundError(f"文件不存在: {input_path}")

    # 读取 Excel 文件
    try:
        df = pd.read_excel(input_path, dtype=str)
    except Exception as e:
        raise ValueError(f"读取 Excel 失败: {e}")

    # 如果指定了筛选表头，进行筛选
    if header_name and filter_values:
        if header_name not in df.columns:
            raise ValueError(f"筛选表头 '{header_name}' 在文件中不存在")

        # 解析筛选值
        values_list = [v.strip() for v in filter_values.split(",") if v.strip()]
        if not values_list:
            raise ValueError("筛选值不能为空")

        # 筛选数据
        filtered_df = df[df[header_name].isin(values_list)].copy()
        if len(filtered_df) == 0:
            raise ValueError(
                f"未找到表头 '{header_name}' 值为 '{filter_values}' 的记录"
            )
    else:
        # 不进行筛选，使用全部数据
        filtered_df = df.copy()

    # 如果指定了要保留的列，只保留这些列
    if columns_to_keep:
        # 检查列是否存在
        available_columns = [
            col for col in columns_to_keep if col in filtered_df.columns
        ]
        missing_columns = [
            col for col in columns_to_keep if col not in filtered_df.columns
        ]

        if missing_columns:
            print(f"警告: 以下列不存在，将被忽略: {missing_columns}")

        if available_columns:
            filtered_df = filtered_df[available_columns]
        else:
            print("警告: 没有找到任何指定的列，将保留所有列")

    # 随机抽样
    if len(filtered_df) <= total_count:
        # 如果数据量小于等于要求数量，使用全部数据
        sampled_df = filtered_df.copy()
        print(
            f"警告: 数据量({len(filtered_df)})小于要求数量({total_count})，使用全部数据"
        )
    else:
        # 随机抽样
        sampled_df = filtered_df.sample(n=total_count, random_state=42)

    # 重新设置序号列
    sampled_df = sampled_df.reset_index(drop=True)
    sampled_df.insert(0, "序号", sampled_df.index + 1)

    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)
    output_path = os.path.join(output_dir, output_file)

    # 保存结果并美化表格
    try:
        # 先保存数据
        sampled_df.to_excel(output_path, index=False, sheet_name="抽查结果")

        # 然后打开文件设置样式
        workbook = openpyxl.load_workbook(output_path)
        worksheet = workbook["抽查结果"]

        # 设置行高
        for row in range(2, len(sampled_df) + 2):  # 从第2行开始（跳过表头）
            worksheet.row_dimensions[row].height = 25

        # 设置表头行高和样式
        header_row = 1
        worksheet.row_dimensions[header_row].height = 30

        # 设置列宽自适应
        for idx, col in enumerate(sampled_df.columns, 1):
            column_letter = get_column_letter(idx)

            # 计算列宽（基于内容长度）
            max_length = 0
            column = sampled_df[col].astype(str)

            # 计算表头宽度
            header_length = len(str(col))
            max_length = max(max_length, header_length)

            # 计算内容最大宽度
            content_length = column.str.len().max()
            max_length = max(max_length, content_length)

            # 设置列宽（添加一些边距）
            adjusted_width = min(max_length + 4, 50)  # 最大宽度限制为50
            worksheet.column_dimensions[column_letter].width = adjusted_width

        # 创建样式对象
        center_alignment = Alignment(
            horizontal="center", vertical="center", wrap_text=True
        )
        bold_font = Font(bold=True, size=12)
        header_fill = PatternFill(
            start_color="D3D3D3", end_color="D3D3D3", fill_type="solid"
        )
        thin_border = Border(
            left=Side(style="thin"),
            right=Side(style="thin"),
            top=Side(style="thin"),
            bottom=Side(style="thin"),
        )

        # 设置所有单元格居中对齐和边框
        for row in worksheet.iter_rows(
            min_row=1,
            max_row=len(sampled_df) + 1,
            min_col=1,
            max_col=len(sampled_df.columns),
        ):
            for cell in row:
                cell.alignment = center_alignment
                cell.border = thin_border

        # 设置表头样式
        for cell in worksheet[1]:
            cell.font = bold_font
            cell.fill = header_fill

        # 保存修改后的文件
        workbook.save(output_path)

    except Exception as e:
        raise ValueError(f"保存结果文件失败: {e}")

    print(f"抽样完成: 共抽取 {len(sampled_df)} 条记录")
    print(f"输出文件: {output_path}")

    return output_path
